"""A tiny shim of langgraph.graph to support tests without external dependency.

This implements a minimal StateGraph that stores node functions and allows invoking
the single node corresponding to state['current_phase']. It also supports
registering a single conditional edge set for 'assess' used in tests.
"""
from typing import Any, Callable, Dict, Optional

END = "__END__"


class CompiledGraph:
    def __init__(self, nodes: Dict[str, Callable], conditional: Dict[str, Any], edges: Dict[str, str], entry_point: Optional[str]):
        self._nodes = nodes
        self._conditional = conditional
        self._edges = edges
        self._entry_point = entry_point

    def get_initial_state(self) -> Dict[str, Any]:
        # Return a minimal initial state expected by tests
        return {"current_phase": self._entry_point}

    def invoke(self, state: Dict[str, Any]) -> Dict[str, Any]:
        # Run nodes in sequence until END or no further transitions.
        current = state.get("current_phase", self._entry_point)
        if current is None:
            return state

        next_state = state
        steps = 0
        max_steps = 50
        while current is not None and current != END and steps < max_steps:
            steps += 1
            # remember previous phase to resolve outgoing edges from it
            prev = current
            node = self._nodes.get(current)
            if node:
                try:
                    next_state = node(next_state)
                except Exception:
                    # if node raises, propagate the state with error
                    raise
            else:
                # no node to handle current phase
                break

            # if node explicitly set a new current_phase, follow it immediately
            new_current = next_state.get("current_phase")
            if new_current is not None and new_current != prev:
                current = new_current
                continue

            # conditional edges (registered via add_conditional_edges) keyed by prev
            if prev in self._conditional:
                try:
                    cond_func = self._conditional[prev]["func"]
                    mapping = self._conditional[prev]["mapping"]
                    key = cond_func(next_state)
                    dest = mapping.get(key)
                    if dest:
                        next_state = {**next_state, "current_phase": dest}
                        current = dest
                        continue
                except Exception:
                    pass

            # fixed edges keyed by prev
            dest = self._edges.get(prev)
            if dest:
                next_state = {**next_state, "current_phase": dest}
                current = dest
                continue

            # no further transitions
            break

        return next_state

    def get_graph(self):
        # Provide minimal object with draw_mermaid used in original code
        class G:
            def draw_mermaid(self):
                return "graph TD; assess-->reactive; assess-->collect_data; collect_data-->analyze; analyze-->recommend; recommend-->respond;"

        return G()


class StateGraph:
    def __init__(self, _state_type=None):
        self._nodes: Dict[str, Callable] = {}
        self._entry_point: Optional[str] = None
        self._conditional: Dict[str, Any] = {}
        self._edges: Dict[str, str] = {}

    def add_node(self, name: str, func: Callable):
        self._nodes[name] = func

    def set_entry_point(self, name: str):
        self._entry_point = name

    def add_conditional_edges(self, name: str, func: Callable, mapping: Dict[str, str]):
        # store function and mapping keyed by name
        self._conditional[name] = {"func": func, "mapping": mapping}

    def add_edge(self, a: str, b: str):
        self._edges[a] = b

    def compile(self) -> CompiledGraph:
        return CompiledGraph(self._nodes, self._conditional, self._edges, self._entry_point)
